package com.hyq.learn.api.config.datasource;

import com.hyq.learn.api.enumtype.DatasourceType;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author heyuqin
 * @date 2024/5/29 15:15
 */
@Configuration
public class DataSourceConfig {

    @Bean("dynamicDataSource")
    public DynamicDataSource dynamicDataSource(BeanFactory beanFactory, DbProperties dbProperties) {
        DefaultListableBeanFactory defaultListableBeanFactory = (DefaultListableBeanFactory) beanFactory;
        Map<DatasourceType, DataSource> datasourceMap = new HashMap<>(dbProperties.getMultiDs().size());
        Map<DatasourceType, DbProperties.LearnDs> multiDs = dbProperties.getMultiDs();

        DbProperties.LearnDs masterDatasource = Optional.ofNullable(multiDs.get(DatasourceType.MASTER))
                .orElseThrow(() -> new RuntimeException("datasource master can not be null"));
        DataSource masterDataSource = buildDataSource(masterDatasource);
        BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(DataSource.class, () -> masterDataSource);
        BeanDefinition beanDefinition = builder.getRawBeanDefinition();
        beanDefinition.setPrimary(true);
        defaultListableBeanFactory.registerBeanDefinition(DatasourceType.MASTER.getValue(), beanDefinition);
        datasourceMap.put(DatasourceType.MASTER, masterDataSource);

        Set<Map.Entry<DatasourceType, DbProperties.LearnDs>> databaseConfigs = multiDs.entrySet().stream()
                .filter(entry -> !DatasourceType.MASTER.equals(entry.getKey())).collect(Collectors.toSet());
        for (Map.Entry<DatasourceType, DbProperties.LearnDs> databaseConfig : databaseConfigs) {
            DataSource dataSource = buildDataSource(databaseConfig.getValue());
            defaultListableBeanFactory.registerSingleton(databaseConfig.getKey().getValue(), dataSource);
            datasourceMap.put(databaseConfig.getKey(), dataSource);
        }
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        dynamicDataSource.setDefaultTargetDataSource(datasourceMap.get(DatasourceType.MASTER));
        Map<Object, Object> targetDataSources = new HashMap<>(datasourceMap);
        dynamicDataSource.setTargetDataSources(targetDataSources);

        return dynamicDataSource;
    }

    /**
     * 事务管理器
     */
    @Bean("dynamicTransactionManager")
    public DataSourceTransactionManager dynamicTransactionManager(@Qualifier("dynamicDataSource") DataSource dataSource) {
        return new DataSourceTransactionManager(dataSource);
    }

    private DataSource buildDataSource(DbProperties.LearnDs learnDs) {
        DataSourceProperties dataSourceProperties = learnDs.getDatasource();
        HikariConfig hikariConfig = learnDs.getHikari();
        hikariConfig.setUsername(dataSourceProperties.getUsername());
        hikariConfig.setPassword(dataSourceProperties.getPassword());
        hikariConfig.setJdbcUrl(dataSourceProperties.getUrl());
        hikariConfig.setDriverClassName(dataSourceProperties.getDriverClassName());
        return new HikariDataSource(hikariConfig);
    }
}
